#pragma once
#include <iostream>
#include <iterator>
#include <map>
#include <regex>
#include <string>
#include <sstream>
#include <unordered_map>
#include <utf8proc.h>

namespace bart {
using wstr_pair = std::pair<std::wstring, std::wstring>;
using Vocab = std::unordered_map<std::wstring, size_t>;
// define a tokenizer only for Bart LLM
class BartTokenizer {
  std::string pat =
      "'s|'t|'re|'ve|'m|'ll|'d| ?[a-zA-Z]+| ?[0-9]+| ?[\\W]+|\\s+(?!\\S)|\\s+";
  std::map<int, char> byte_encoder;
  std::map<wstr_pair, int> bpe_rank;
  std::shared_ptr<Vocab> vocab;

  std::vector<std::wstring> tokenize(const std::string &str);
  wstr_pair get_min(const std::vector<wstr_pair> pairs);
  std::wstring bpe(const std::wstring &token);
  std::vector<int> convert_to_ids(const std::vector<std::wstring> &tokens);

public:
  BartTokenizer(const std::string &dir_path = "../model/");
  std::vector<int> encode(const std::string &str, const std::string &text_pair);
};

void printEncoding(const std::vector<int> &ids);
} // namespace bart